Skip to content

Conversation

@jpienaar
Copy link
Member

Previously this would assert when attempting to getMutableData.

Previously this would assert when attempting to getMutableData.
@jpienaar jpienaar requested a review from River707 September 14, 2024 04:10
@llvmbot llvmbot added mlir:core MLIR Core Infrastructure mlir labels Sep 14, 2024
@llvmbot
Copy link
Member

llvmbot commented Sep 14, 2024

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Jacques Pienaar (jpienaar)

Changes

Previously this would assert when attempting to getMutableData.


Full diff: https://github.com/llvm/llvm-project/pull/108679.diff

2 Files Affected:

  • (modified) mlir/include/mlir/IR/AsmState.h (+8-1)
  • (modified) mlir/unittests/IR/AttributeTest.cpp (+15)
diff --git a/mlir/include/mlir/IR/AsmState.h b/mlir/include/mlir/IR/AsmState.h
index 42cbedcf9f8837..edbd3bb6fc15db 100644
--- a/mlir/include/mlir/IR/AsmState.h
+++ b/mlir/include/mlir/IR/AsmState.h
@@ -82,6 +82,8 @@ class AsmStateImpl;
 //===----------------------------------------------------------------------===//
 // Resource Entry
 
+class HeapAsmResourceBlob;
+
 /// This class represents a processed binary blob of data. A resource blob is
 /// essentially a collection of data, potentially mutable, with an associated
 /// deleter function (used if the data needs to be destroyed).
@@ -177,6 +179,8 @@ class AsmResourceBlob {
 
   /// Whether the data is mutable.
   bool dataIsMutable;
+
+  friend class HeapAsmResourceBlob;
 };
 
 /// This class provides a simple utility wrapper for creating heap allocated
@@ -196,8 +200,11 @@ class HeapAsmResourceBlob {
   static AsmResourceBlob allocateAndCopyWithAlign(ArrayRef<char> data,
                                                   size_t align,
                                                   bool dataIsMutable = true) {
-    AsmResourceBlob blob = allocate(data.size(), align, dataIsMutable);
+    // This sets the blob to be mutable initially to allow writing
+    // (getMutableData) below.
+    AsmResourceBlob blob = allocate(data.size(), align, /*dataIsMutable=*/true);
     std::memcpy(blob.getMutableData().data(), data.data(), data.size());
+    blob.dataIsMutable = dataIsMutable;
     return blob;
   }
   template <typename T>
diff --git a/mlir/unittests/IR/AttributeTest.cpp b/mlir/unittests/IR/AttributeTest.cpp
index e72bfe9d82e7cf..5a3649d7be923a 100644
--- a/mlir/unittests/IR/AttributeTest.cpp
+++ b/mlir/unittests/IR/AttributeTest.cpp
@@ -351,6 +351,21 @@ TEST(DenseResourceElementsAttrTest, CheckNoCast) {
   EXPECT_FALSE(isa<DenseBoolResourceElementsAttr>(i32ResourceAttr));
 }
 
+TEST(DenseResourceElementsAttrTest, CheckNotMutableAllocateAndCopy) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  // Create a i32 attribute.
+  std::vector<int32_t> data = {10, 20, 30};
+  auto type = RankedTensorType::get(data.size(), builder.getI32Type());
+  Attribute i32ResourceAttr = DenseI32ResourceElementsAttr::get(
+      type, "resource",
+      HeapAsmResourceBlob::allocateAndCopyInferAlign<int32_t>(
+          data, /*is_mutable=*/false));
+
+  EXPECT_TRUE(isa<DenseI32ResourceElementsAttr>(i32ResourceAttr));
+}
+
 TEST(DenseResourceElementsAttrTest, CheckInvalidData) {
   MLIRContext context;
   Builder builder(&context);

@jpienaar jpienaar merged commit b96ebee into llvm:main Oct 11, 2024
DanielCChen pushed a commit to DanielCChen/llvm-project that referenced this pull request Oct 16, 2024
Previously this would assert when attempting to getMutableData.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

mlir:core MLIR Core Infrastructure mlir

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants